from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn
import argparse
from faiss_search import FaissSearch
from fastapi import Body, FastAPI, File, Form, Query, UploadFile, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from typing import Union
from io import BytesIO

import pydantic
from utils import *

import logging

# Configure logging
logging.basicConfig(filename='api.log', level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')


class BaseResponse(BaseModel):
    code: int = pydantic.Field(200, description="HTTP status code")
    msg: str = pydantic.Field("success", description="HTTP status message")
    size_before: int = Body(..., description="操作之前数据库大小", example=0)
    size_after: int = Body(..., description="操作之后数据库大小", example=0)

    class Config:
        schema_extra = {
            "example": {
                "code": 200,
                "msg": "success",
                "size_before":0,
                "size_after":1,
            }
        }


#########################################################
########################## Get ##########################
#########################################################

class InfoResponse(BaseModel):
    code: int = 200
    msg: str = "success"
    info: dict

    class Config:
        schema_extra = {
            "example": {
                "code": 200,
                "msg": "success",
                "info": {
                    "vector_store_path": "./Faiss/bge-small/vector_store",
                    "vector_store_type": "FAISS",
                    "vector_store_size": 500,
                    "emb_model_path": "./Embeddings/bge-small-zh-v1.5",
                    "emb_model": "bge-small",
                    "emb_dim": 512
                    }
                }
        }


async def get_info():
    try:
        info = vanilla_faiss.get_info()
        return InfoResponse(code=200, msg="success", info=info)
    except Exception as e:
        msg = f"fail: {e}"
        return InfoResponse(code=500, msg=msg, info={})

async def keyword_get_info():
    try:
        info = keyword_faiss.get_info()
        return InfoResponse(code=200, msg="success", info=info)
    except Exception as e:
        msg = f"fail: {e}"
        return InfoResponse(code=500, msg=msg, info={})

async def mysql_get_info():
    try:
        info = mysql_faiss.get_info()
        return InfoResponse(code=200, msg="success", info=info)
    except Exception as e:
        msg = f"fail: {e}"
        return InfoResponse(code=500, msg=msg, info={})



#########################################################
########################## Add ##########################
#########################################################


async def add_from_excel(excel_file: UploadFile):
    size_before = vanilla_faiss.get_size()
    try:
        content = await excel_file.read()
        content_io = BytesIO(content)
        texts, metadata = read_excel_file(content_io, excel_file.filename)
        vanilla_faiss.add_texts(texts, metadata)
        size_after = vanilla_faiss.get_size()
        n_added = size_after - size_before
        msg = f"success: {n_added}条数据成功添加"
        return BaseResponse(code=200, msg=msg, size_before=size_before, size_after=size_after)
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)


async def keyword_add_from_excel(excel_file: UploadFile):
    size_before = keyword_faiss.get_size()
    try:
        content = await excel_file.read()
        content_io = BytesIO(content)
        texts, metadata = keyword_read_excel_file(content_io, excel_file.filename)
        keyword_faiss.add_texts(texts, metadata)
        size_after = keyword_faiss.get_size()
        n_added = size_after - size_before
        msg = f"success: {n_added}条数据成功添加"
        return BaseResponse(code=200, msg=msg, size_before=size_before, size_after=size_after)
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)


async def add_from_json(json_file: UploadFile):
    size_before = vanilla_faiss.get_size()
    try:
        content = await json_file.read()
        json_data = json.loads(content)
        texts, metadatas = read_json_file(json_data, json_file.filename)
        
        vanilla_faiss.add_texts(texts, metadatas)
        size_after = vanilla_faiss.get_size()
        n_added = size_after - size_before
        msg = f"success: {n_added}条数据成功添加"
        return BaseResponse(code=200, msg=msg, size_before=size_before, size_after=size_after)
    
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)

async def add_from_MySQL():
    size_before = mysql_faiss.get_size()
    try:
        # content = await json_file.read()
        # json_data = json.loads(content)
        # texts, metadata = read_json_file(json_data, json_file.filename)
        last_update_date = mysql_faiss.last_update_date
        texts, metadatas, last_update_date = read_MySQL_DB(min_time=last_update_date)
        print("last_update_data: ", last_update_date)
        mysql_faiss.add_texts(texts, metadatas)
        
        print("the params: ", mysql_faiss.last_update_date)
        if last_update_date:
            mysql_faiss.last_update_date = last_update_date
            mysql_faiss.save_params()
        size_after = mysql_faiss.get_size()
        n_added = size_after - size_before
        msg = f"success: {n_added}条数据成功添加"
        return BaseResponse(code=200, msg=msg, size_before=size_before, size_after=size_after)
    
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)




async def keyword_add_from_json(json_file: UploadFile):
    size_before = keyword_faiss.get_size()
    try:
        content = await json_file.read()
        json_data = json.loads(content)
        texts, metadata = keyword_read_json_file(json_data, json_file.filename)
        keyword_faiss.add_texts(texts, metadata)

        size_after = keyword_faiss.get_size()
        n_added = size_after - size_before
        msg = f"success: {n_added}条数据成功添加"
        return BaseResponse(code=200, msg=msg, size_before=size_before, size_after=size_after)
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)


#########################################################
######################### Search ########################
#########################################################
class SearchRequest(BaseModel):
    query_text: str = Body(..., description="问题", example="工伤保险办理")
    top_k: int = Body(None, description="返回结果数量", example=5)
    threshold: float = Body(0.75, description="相似度阈值", example=50.0)
    class Config:
        schema_extra = {
            "example": {
                "query_text": "企业创业补贴如何办理",
                "top_k": 5,
                "threshold":50.0
            }
        }


class SearchResponse(BaseModel):
    code: int = 200
    msg: str = "success"
    ids: list
    scores:list
    texts: list
    # similarities: list
    # results: dict
    results_before_filter: int
    results_after_filter: int
    
    class Config:
        schema_extra = {
            "example": {
                "code": 200,
                "msg": "success",
                "results":{
                    '市民来电咨询:企业在一网通办惠企点办理企业创业补贴贷款，提交材料时提示资料提交失败，业务办理部门：人力资源。': 25.472061,
                    '市民来电咨询:其为内资企业，有一个董事代表要离职。现咨询：需要办理什么手续和在哪里办。': 40.604416,
                    '市民来电反复催单。催单原因：市民称管理部门至今未给予任何回复，请管理部门尽快处理。【最近派发的工单编号：20230201012753，工单内容：市民来电反映:其分公司是一家有限公司，注册在闵行区，现在想变更分公司企业法人，询问具体流程和材料，望管理部门指导告知如何操作。（需回复）】': 45.609825,
                    '市民来电咨询:如何办理 对外贸易经营者备案登记': 50.554813,
                    '市民来电反映:2023年3月9日上传了工商年报，从国家企业信息系统上无法查看到信息，希望管理部门能尽快核实更新信息。': 51.99854
                    },
                "results_before_filter":10,
                "results_after_filter":8,
                }
        }


async def similarity_search(search_request: SearchRequest):
    query_text = search_request.query_text
    top_k = search_request.top_k
    if top_k is None:
        # top_k = mysql_faiss.get_size()
        top_k = vanilla_faiss.get_size()
    threshold = search_request.threshold
    ids = []
    try:
        # results = mysql_faiss.similarity_search(query_text, top_k)
        results = vanilla_faiss.similarity_search(query_text, top_k)
        results_before_filter = len(results)
        if threshold:
            results = {key: value for key, value in results.items() if (2-value['score'])/2 >= threshold}
        ids = [value["id"] if "id" in value else None for value in results.values() ]
        scores = [value["score"] if "score" in value else None for value in results.values()]
        scores = [(2-i)/2 for i in scores]
        texts = [key for key, value in results.items()]
        results_after_filter = len(results)
        # vector_store_size = mysql_faiss.get_size()
        vector_store_size = vanilla_faiss.get_size()
        msg = f"success: 从{vector_store_size}条数据中, {results_after_filter}条数据成功返回"
        if results_before_filter != top_k:
            n_filtered = top_k - results_after_filter
            msg += f"（{n_filtered}条重复或不相似数据被筛除或不存在）"
        # ids = [{"id":i, "score":j} for i,j in zip(ids,scores)]
        ids = [{"id":i, "score":j,"content":z} for i,j,z in zip(ids,scores,texts)]
        # return SearchResponse(code=200, msg=msg, ids=ids, scores=scores, texts=texts, results=results,results_before_filter=results_before_filter,results_after_filter=results_after_filter)
        return SearchResponse(code=200, msg=msg, ids=ids, scores=scores, texts=texts,results_before_filter=results_before_filter,results_after_filter=results_after_filter)
    
    except Exception as e:
        msg = f"fail: {e}"
        return SearchResponse(code=500, msg=msg, ids=[], scores=[], texts=[], results={},results_before_filter=0,results_after_filter=0)


async def mysql_similarity_search(search_request: SearchRequest):
    query_text = search_request.query_text
    top_k = search_request.top_k
    if top_k is None:
        top_k = mysql_faiss.get_size()
    threshold = search_request.threshold
    ids = []
    try:
        results = mysql_faiss.similarity_search(query_text, top_k)
        results_before_filter = len(results)
        if threshold:
            results = {key: value for key, value in results.items() if (2-value['score'])/2 >= threshold}
        ids = [value["id"] if "id" in value else None for value in results.values() ]
        scores = [value["score"] if "score" in value else None for value in results.values()]
        scores = [(2-i)/2 for i in scores]
        texts = [key for key, value in results.items()]
        results_after_filter = len(results)
        vector_store_size = mysql_faiss.get_size()
        msg = f"success: 从{vector_store_size}条数据中, {results_after_filter}条数据成功返回"
        if results_before_filter != top_k:
            n_filtered = top_k - results_after_filter
            msg += f"（{n_filtered}条重复或不相似数据被筛除或不存在）"
        ids = [{"id":i, "score":j} for i,j in zip(ids,scores)]
        return SearchResponse(code=200, msg=msg, ids=ids, scores=scores, texts=texts, results=results,results_before_filter=results_before_filter,results_after_filter=results_after_filter)
    except Exception as e:
        msg = f"fail: {e}"
        return SearchResponse(code=500, msg=msg, ids=[], scores=[], texts=[], results={},results_before_filter=0,results_after_filter=0)


async def keyword_similarity_search(search_request: SearchRequest):
    query_text = search_request.query_text
    top_k = search_request.top_k
    if top_k is None:
        top_k = keyword_faiss.get_size()
    threshold = search_request.threshold
    ids = []
    try:
        results = keyword_faiss.similarity_search(query_text, top_k)
        
        # ids = [value["id"] if "id" in value else None for value in results.values() ]

        results_before_filter = len(results)
        if threshold:
            results = {key: value for key, value in results.items() if (2-value['score'])/2 >= threshold}
        ids = [value["id"] if "id" in value else None for value in results.values() ]
        scores = [value["score"] if "score" in value else None for value in results.values() ]
        scores = [(2-i)/2 for i in scores]
        texts = [key for key, value in results.items()]
        results_after_filter = len(results)

        vector_store_size = keyword_faiss.get_size()
        msg = f"success: 从{vector_store_size}条数据中, {results_after_filter}条数据成功返回"
        if results_before_filter != top_k:
            n_filtered = top_k - results_after_filter
            msg += f"（{n_filtered}条重复或不相似数据被筛除或不存在）"
        ids = [{"id":i, "score":j} for i,j in zip(ids,scores)]
        return SearchResponse(code=200, msg=msg,ids=ids, scores=scores, texts=texts, results=results,results_before_filter=results_before_filter,results_after_filter=results_after_filter)
    except Exception as e:
        msg = f"fail: {e}"
        return SearchResponse(code=500, msg=msg, ids=[],scores=[], texts=[],results={},results_before_filter=0,results_after_filter=0)


#########################################################
######################### Delete ########################
#########################################################


async def delete_all():
    size_before = vanilla_faiss.get_size()
    try:
        vanilla_faiss.delete_all()
        size_after = vanilla_faiss.get_size()
        msg = f"success: {size_before}条数据成功删除"
        return BaseResponse(code=200, msg=msg,size_before=size_before, size_after=size_after)
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)


async def keyword_delete_all():
    size_before = keyword_faiss.get_size()
    try:
        keyword_faiss.delete_all()
        size_after = keyword_faiss.get_size()
        msg = f"success: {size_before}条数据成功删除"
        return BaseResponse(code=200, msg=msg,size_before=size_before, size_after=size_after)
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)


async def mysql_delete_all():
    size_before = mysql_faiss.get_size()
    try:
        mysql_faiss.delete_all()
        size_after = mysql_faiss.get_size()
        msg = f"success: {size_before}条数据成功删除"
        return BaseResponse(code=200, msg=msg,size_before=size_before, size_after=size_after)
    except Exception as e:
        msg = f"fail: {e}"
        return BaseResponse(code=500, msg=msg,size_before=size_before, size_after=size_before)






def api_start(host, port):
    global vanilla_faiss
    global keyword_faiss
    global mysql_faiss
    global app

    app = FastAPI()

    # 
    #     app.add_middleware(
    #         CORSMiddleware,
    #         allow_origins=["*"],
    #         allow_credentials=True,
    #         allow_methods=["*"],
    #         allow_headers=["*"],
    #     )

    app.get("/faiss/get_info", response_model=InfoResponse)(get_info)
    app.get("/faiss/keyword_get_info", response_model=InfoResponse)(keyword_get_info)
    app.get("/faiss/mysql_get_info", response_model=InfoResponse)(mysql_get_info)
    

    
    app.post("/faiss/add_from_json", response_model=BaseResponse)(add_from_json)
    app.post("/faiss/add_from_excel", response_model=BaseResponse)(add_from_excel)

    app.post("/faiss/keyword_add_from_excel", response_model=BaseResponse)(keyword_add_from_excel)
    app.post("/faiss/keyword_add_from_json", response_model=BaseResponse)(keyword_add_from_json)

    
    app.post("/faiss/add_from_MySQL", response_model=BaseResponse)(add_from_MySQL)

    app.post("/faiss/similarity_search", response_model=SearchResponse)(similarity_search)
    app.post("/faiss/mysql_similarity_search", response_model=SearchResponse)(mysql_similarity_search)
    app.post("/faiss/keyword_similarity_search", response_model=SearchResponse)(keyword_similarity_search)
    
    app.delete("/faiss/delete_all", response_model=BaseResponse)(delete_all)
    app.delete("/faiss/keyword_delete_all", response_model=BaseResponse)(keyword_delete_all)
    app.delete("/faiss/mysql_delete_all", response_model=BaseResponse)(mysql_delete_all)



    vanilla_faiss = FaissSearch()
    vanilla_faiss.load()

    keyword_faiss = FaissSearch(type="keyword")
    keyword_faiss.load()

    mysql_faiss = FaissSearch(type="mysql")
    mysql_faiss.load()

    uvicorn.run(app, host=host, port=port)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="使用FAISS查找相似文本")
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7999)


    # 初始化消息
    args = None
    args = parser.parse_args()
    api_start(args.host, args.port)
    
    

# uvicorn api2:app --host 0.0.0.0 --port 8000 --reload
# POST: http://10.4.251.99:8000/similarity_search
# http://10.4.251.99:8000/docs


